# @Author ltj
# @Time   2021/11/8 14:04
# @File   Information.py

import qrcode
from blind_watermark import WaterMark
from PIL import Image
import tranferToPic
import socket
import time
import os
import random
import string
from pyzbar import pyzbar
from ocr_onnx.text_recognizer import TextRecognizer
import cv2
import numpy as np

# 图片存在example\output里
OUTPUT_FOLD = "../examples/output/"


# 用户请求图片时读取用户信息生成二维码嵌入图片
def getPicture(IP, userName, loadTime):
    qr = qrcode.QRCode(
        version=1,
        error_correction=qrcode.constants.ERROR_CORRECT_Q,
        box_size=2,
        border=1,
    )
    data = 'username: ' + userName + ' IP: ' + IP + ' loadTime: ' + loadTime
    img_file = r'..\examples\output\py_qrcode_' + userName + '.png'
    qr.add_data(data)

    qr.make(fit=True)
    img = qr.make_image(fill_color="black")
    img = img.resize((150, 150), Image.ANTIALIAS)
    img.save(img_file)

    bwm1 = WaterMark(password_wm=1, password_img=1)
    # 读取原图
    bwm1.read_img('../examples/pic/ori_img.jpg')
    # 读取水印
    bwm1.read_wm('../examples/output/py_qrcode_' + userName + '.png')
    # 打上盲水印
    bwm1.embed('../examples/output/embed_img_' + userName + '.png')


# 用户请求读取图片身份信息时返回包含身份信息的二维码图片
def getInfo(pictureName):
    # image = Image.open('../examples/output/py_qrcode_' + pictureName + '.png')
    # imagePixmap = image.size
    # print(imagePixmap)
    bwm1 = WaterMark(password_wm=1, password_img=1)
    # 注意需要设定水印的长宽wm_shape
    bwm1.extract(filename='../examples/output/embed_img_' + pictureName + '.png', wm_shape=(150, 150),
                 out_wm_name='../examples/output'
                             '/extracted_img_' + pictureName + '.png', )


# 读取用户信息
def getIP():
    hostname = socket.gethostname()  # 获取本机IP
    user_name = socket.gethostname()  # 获取用户名
    lode_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())  # 获取下载时间
    return hostname, user_name, lode_time


# 解析二维码
def decode_wm(pictureName):
    code_img_path = '../examples/output/extracted_img_' + pictureName + '.png'
    if not os.path.exists(code_img_path):
        raise FileExistsError(code_img_path)
    return pyzbar.decode(Image.open(code_img_path))[0].data.decode("utf-8")


# 把str嵌入到图片中
def embedded_by_image_and_str(img_name, str):
    bwm1 = WaterMark(password_wm=1, password_img=1)

    # 读取原图
    bwm1.read_img(OUTPUT_FOLD + img_name)

    # 根据str获得要嵌入的水印
    mark_name = tranferToPic.tranferToPic(str)

    # 读取水印
    bwm1.read_wm(OUTPUT_FOLD + mark_name)

    # 打上盲水印
    embedded_name = ''.join(random.sample(string.ascii_letters + string.digits, 16)) + '.png'
    bwm1.embed(OUTPUT_FOLD + embedded_name)

    return embedded_name


# 提取图片中的水印，识别成string返回
def getStrByImage(img_name):
    # 提取水印
    bwm1 = WaterMark(password_wm=1, password_img=1)
    embedded_name = ''.join(random.sample(string.ascii_letters + string.digits, 16)) + '.png'
    bwm1.extract(filename=OUTPUT_FOLD + img_name, wm_shape=(150, 150),
                 out_wm_name=OUTPUT_FOLD + embedded_name, )

    #文本识别
    result = identify_watermark(OUTPUT_FOLD + embedded_name)
    if result == -1:
        return "-1"
    else:
        return result.replace(' ', '')


def identify_watermark(file_path):
    model_path = os.path.join('..', 'ocr_onnx', 'ocr', 'ocr.onnx')
    char_dict_path = os.path.join('..', 'ocr_onnx', 'ocr', 'key.txt')
    text_recognizer = TextRecognizer(model_path, char_dict_path)
    image = cv2.imread(file_path)
    sub1 = image[0:50, 0:150]
    sub2 = image[50:100, 0:150]
    sub3 = image[100:150, 0:150]
    final_matrix = np.zeros((50, 450, 3), np.uint8)
    # change
    final_matrix[0:50, 0:150] = sub1
    final_matrix[0:50, 150:300] = sub2
    final_matrix[0:50, 300:450] = sub3
    # image = image[:30, 977:1058, :]
    text, score = text_recognizer.process(final_matrix)
    if score > 0.5:
        return text
    else:
        return -1


if __name__ == "__main__":
    identify_watermark(r"C:\Users\ASUS\Desktop\1.png")
    # arr = getIP()
    # IP = arr[0]
    # userName = arr[1]
    # loadTime = arr[2]
    # getPicture(IP, userName, loadTime)
    # getInfo(userName)
    # print(decode_wm(userName))
